#include "SkinnedMesh.h"

class BONE_HIERARCHY: public ID3DXAllocateHierarchy
{
public:
	STDMETHOD(CreateFrame)(THIS_ LPCSTR Name, LPD3DXFRAME *ppNewFrame);
	STDMETHOD(CreateMeshContainer)(THIS_ LPCTSTR Name, CONST D3DXMESHDATA * pMeshData, CONST D3DXMATERIAL * pMaterials, CONST D3DXEFFECTINSTANCE * pEffectInstances, DWORD NumMaterials, CONST DWORD * pAdjacency, LPD3DXSKININFO pSkinInfo, LPD3DXMESHCONTAINER * ppNewMeshContainer);
	STDMETHOD(DestroyFrame)(THIS_ LPD3DXFRAME pFrameToFree);
	STDMETHOD(DestroyMeshContainer)(THIS_ LPD3DXMESHCONTAINER pMeshContainerBase);
};

HRESULT BONE_HIERARCHY::CreateFrame(LPCSTR Name, LPD3DXFRAME* ppNewFrame)
{
	BONE* pNewBone = new BONE;
	memset(pNewBone, 0, sizeof(BONE));

	//	Copy the name over
	if(Name != NULL)
	{
		pNewBone->Name = new char[strlen(Name) + 1];
		strcpy(pNewBone->Name, Name);
	}
	//	Set the transformation matrices
	D3DXMatrixIdentity(&pNewBone->TransformationMatrix);
	D3DXMatrixIdentity(&pNewBone->CombinedTransformationMatrix);

	//	return the newly created bone
	*ppNewFrame = (D3DXFRAME*)pNewBone;

	return S_OK;
}

HRESULT BONE_HIERARCHY::CreateMeshContainer(LPCSTR Name,
											CONST D3DXMESHDATA *pMeshData,
											CONST D3DXMATERIAL *pMaterials,
											CONST D3DXEFFECTINSTANCE *pEffectInstances,
											DWORD NumMaterials,
											CONST DWORD *pAdjacency,
											LPD3DXSKININFO pSkinInfo,
											LPD3DXMESHCONTAINER *ppNewMeshContainer)
{

	//	Create new Bone Mesh
	BONEMESH* pBoneMesh = new BONEMESH;
	memset(pBoneMesh, 0, sizeof(BONEMESH));

	//	Get the mesh data from the X file
	pBoneMesh->OriginalMesh = pMeshData->pMesh;
	pBoneMesh->MeshData.pMesh = pMeshData->pMesh;
	pBoneMesh->MeshData.Type = pMeshData->Type;
	pMeshData->pMesh->AddRef();			//	Add a reference so that the mesh isn't deallocated
	IDirect3DDevice9* m_pDevice = NULL;
	pMeshData->pMesh->GetDevice(&m_pDevice);		//	Get device pointer from mesh

	//	Copy materials and load textures
	for(unsigned int nMat = 0; nMat < NumMaterials; nMat++)
	{
		D3DXMATERIAL mtrl;
		memcpy(&mtrl, &pMaterials[nMat], sizeof(D3DXMATERIAL));
		pBoneMesh->m_vMaterials.push_back(mtrl.MatD3D);

		char textureFileName[200];
		strcpy(textureFileName, "mesh/");
		strcat(textureFileName, mtrl.pTextureFilename);

		//	Load the texture
		IDirect3DTexture9* pNewTexture = NULL;
		D3DXCreateTextureFromFile(m_pDevice, textureFileName, &pNewTexture);
		pBoneMesh->m_vTextures.push_back(pNewTexture);
	}

	if(pSkinInfo != NULL)
	{
		//	Get skin info
		pBoneMesh->pSkinInfo = pSkinInfo;
		pSkinInfo->AddRef();			//	Add reference so it is no deallocated

		//	Clone mesh and store it in pBoneMesh->MeshData.pMesh
		pMeshData->pMesh->CloneMeshFVF(D3DXMESH_MANAGED, pMeshData->pMesh->GetFVF(),
									   m_pDevice, &pBoneMesh->MeshData.pMesh);

		//	Get attribute table
		pBoneMesh->MeshData.pMesh->GetAttributeTable(NULL, &pBoneMesh->NumAttributeGroups);
		pBoneMesh->attributeTable = new D3DXATTRIBUTERANGE[pBoneMesh->NumAttributeGroups];
		pBoneMesh->MeshData.pMesh->GetAttributeTable(pBoneMesh->attributeTable, NULL);

		//	Create bone offset and the current matrices
		int nNumBones = pSkinInfo->GetNumBones();
		pBoneMesh->boneOffsetMatrices = new D3DXMATRIX[nNumBones];
		pBoneMesh->currentBoneMatrices = new D3DXMATRIX[nNumBones];

		//	Get bone offset matrices
		for(int nBone = 0; nBone < nNumBones; nBone++)
			pBoneMesh->boneOffsetMatrices[nBone] = *(pBoneMesh->pSkinInfo->GetBoneOffsetMatrix(nBone));
	}
	
	//	Set new mesh container to newly created bonemesh
	*ppNewMeshContainer = pBoneMesh;
	return S_OK;
}

HRESULT BONE_HIERARCHY::DestroyFrame(LPD3DXFRAME pFrameToFree)
{
	if(pFrameToFree)
	{
		if(pFrameToFree->Name != NULL)
			delete[] pFrameToFree->Name;
		delete pFrameToFree;
	}
	pFrameToFree = NULL;

	return S_OK;
}

HRESULT BONE_HIERARCHY::DestroyMeshContainer(LPD3DXMESHCONTAINER pMeshContainerBase)
{
	BONEMESH* pBoneMesh = (BONEMESH*)pMeshContainerBase;

	//	Release textures
	for(size_t tex = 0; tex < pBoneMesh->m_vTextures.size(); tex++)
		if(pBoneMesh->m_vTextures[tex] != NULL)
			pBoneMesh->m_vTextures[tex]->Release();
	
	//	Release mesh data
	if(pBoneMesh->MeshData.pMesh)pBoneMesh->MeshData.pMesh->Release();
	if(pBoneMesh->pSkinInfo)pBoneMesh->pSkinInfo->Release();
	if(pBoneMesh->OriginalMesh)pBoneMesh->OriginalMesh->Release();
	delete pBoneMesh;

	return S_OK;
}

//////////////////////////////////////////////////////////////////////////
//	SKINNED MESH														//
//////////////////////////////////////////////////////////////////////////

struct VERTEX
{
	//	Variables
	D3DXVECTOR3 position;
	D3DCOLOR color;

	VERTEX();
	VERTEX(D3DXVECTOR3 pos, D3DCOLOR col) {position = pos; color = col;}

	static const DWORD FVF;
};

const DWORD VERTEX::FVF = D3DFVF_XYZ | D3DFVF_DIFFUSE;

CSkinnedMesh::CSkinnedMesh(void)
{
	m_pRootBone = NULL;
	m_pSphereMesh = NULL;
	m_pAnimControl = NULL;

}

CSkinnedMesh::~CSkinnedMesh(void)
{
	BONE_HIERARCHY boneHierarchy;
	boneHierarchy.DestroyFrame(m_pRootBone);
	if(m_pAnimControl)m_pAnimControl->Release();
}

void CSkinnedMesh::Load(char* szFilename, IDirect3DDevice9* dev)
{
	m_pDevice = dev;
	
	BONE_HIERARCHY boneHierarchy;

	D3DXLoadMeshHierarchyFromX(szFilename, D3DXMESH_MANAGED, m_pDevice, &boneHierarchy,
							   NULL, &m_pRootBone, &m_pAnimControl);

	SetupBoneMatrixPointers((BONE*)m_pRootBone);
	
	//	Update all the bones
	D3DXMATRIX i;
	D3DXMatrixIdentity(&i);
	UpdateMatrices((BONE*)m_pRootBone, &i);

	//	Create the sphere's representing joints
	D3DXCreateSphere(m_pDevice, 0.07f, 10, 10, &m_pSphereMesh, NULL);
}

void CSkinnedMesh::UpdateMatrices(BONE* bone, D3DXMATRIX* parentMatrix)
{
	if(bone == NULL) return;

	D3DXMatrixMultiply(&bone->CombinedTransformationMatrix,
					   &bone->TransformationMatrix,
					   parentMatrix);

	if(bone->pFrameSibling)UpdateMatrices((BONE*)bone->pFrameSibling, parentMatrix);
	if(bone->pFrameFirstChild)UpdateMatrices((BONE*)bone->pFrameFirstChild, &bone->CombinedTransformationMatrix);
}

void CSkinnedMesh::Render(BONE* bone)
{
	if(bone == NULL)bone = (BONE*)m_pRootBone;

	//	If there is a mesh to render
	if(bone->pMeshContainer != NULL)
	{
		BONEMESH* pBoneMesh = (BONEMESH*)bone->pMeshContainer;

		if(pBoneMesh->pSkinInfo != NULL)
		{
			//	Setup the bone transforms
			int nNumBones = pBoneMesh->pSkinInfo->GetNumBones();
			for(int nBone = 0; nBone < nNumBones; nBone++)
				D3DXMatrixMultiply(&pBoneMesh->currentBoneMatrices[nBone],
								   &pBoneMesh->boneOffsetMatrices[nBone],
								   pBoneMesh->boneMatrixPtrs[nBone]);

			//	Update the skinned mesh
			BYTE* pSrc = NULL, *pDest = NULL;
			pBoneMesh->OriginalMesh->LockVertexBuffer(D3DLOCK_READONLY, (VOID**)&pSrc);
			pBoneMesh->MeshData.pMesh->LockVertexBuffer(0, (VOID**)&pDest);

			pBoneMesh->pSkinInfo->UpdateSkinnedMesh(pBoneMesh->currentBoneMatrices, NULL, pSrc, pDest);

			pBoneMesh->MeshData.pMesh->UnlockVertexBuffer();
			pBoneMesh->OriginalMesh->UnlockVertexBuffer();

			//	Render the mesh
			for(int nAttr = 0; nAttr < pBoneMesh->NumAttributeGroups; nAttr++)
			{
				int nMtrlIndex = pBoneMesh->attributeTable[nAttr].AttribId;
				m_pDevice->SetMaterial(&(pBoneMesh->m_vMaterials[nMtrlIndex]));
				m_pDevice->SetTexture(0, pBoneMesh->m_vTextures[nMtrlIndex]);
				pBoneMesh->MeshData.pMesh->DrawSubset(nMtrlIndex);
			}
		}
	}

	if(bone->pFrameSibling != NULL)Render((BONE*)bone->pFrameSibling);
	if(bone->pFrameFirstChild != NULL)Render((BONE*)bone->pFrameFirstChild);
}

void CSkinnedMesh::RenderSkeleton(BONE* bone, BONE* parent, D3DXMATRIX world)
{
	if(world == NULL)return;
	if(bone == NULL) bone = (BONE*)m_pRootBone;

	D3DXMATRIX r, s;
	D3DXMatrixRotationYawPitchRoll(&r, -D3DX_PI * 0.5f, 0.0f, 0.0f);

	//	Draw the spheres
	m_pDevice->SetRenderState(D3DRS_LIGHTING, true);
	m_pDevice->SetTransform(D3DTS_WORLD, &(r * bone->CombinedTransformationMatrix * world));
	if(bone->Name != NULL) m_pSphereMesh->DrawSubset(0);

	//	Draw the lines between the joints
	if(parent != NULL && bone->Name != NULL && parent->Name != NULL)
	{
		D3DXMATRIX w1 = bone->CombinedTransformationMatrix;
		D3DXMATRIX w2 = parent->CombinedTransformationMatrix;

		//	Extract translation
		D3DXVECTOR3 thisBone = D3DXVECTOR3(w1(3,0), w1(3,1), w1(3,2));
		D3DXVECTOR3 parentBone = D3DXVECTOR3(w2(3,0), w2(3,1), w2(3,2));

		if(D3DXVec3Length(&(thisBone - parentBone)) < 2.0f)
		{
			m_pDevice->SetTransform(D3DTS_WORLD, &world);
			VERTEX vert[] = {VERTEX(parentBone, 0xffff0000), VERTEX(thisBone, 0xff00ff00)};
			m_pDevice->SetRenderState(D3DRS_LIGHTING, false);
			m_pDevice->SetFVF(VERTEX::FVF);
			m_pDevice->DrawPrimitiveUP(D3DPT_LINESTRIP, 1, &vert[0], sizeof(VERTEX));
		}
	}

	if(bone->pFrameSibling)RenderSkeleton((BONE*)bone->pFrameSibling, parent, world);
	if(bone->pFrameFirstChild)RenderSkeleton((BONE*)bone->pFrameFirstChild, bone, world);
}

void CSkinnedMesh::SetupBoneMatrixPointers(BONE* bone)
{
	if(bone->pMeshContainer != NULL)
	{
		BONEMESH* pBoneMesh = (BONEMESH*)bone->pMeshContainer;

		if(pBoneMesh->pSkinInfo != NULL)
		{
			int nNumBones = pBoneMesh->pSkinInfo->GetNumBones();
			pBoneMesh->boneMatrixPtrs = new D3DXMATRIX*[nNumBones];

			for(int nBone = 0; nBone < nNumBones; nBone++)
			{
				BONE* b = (BONE*)D3DXFrameFind(m_pRootBone, pBoneMesh->pSkinInfo->GetBoneName(nBone));
				if(b != NULL)pBoneMesh->boneMatrixPtrs[nBone] = &b->CombinedTransformationMatrix;
				else pBoneMesh->boneMatrixPtrs[nBone] = NULL;
			}
		}
	}

	if(bone->pFrameSibling != NULL)SetupBoneMatrixPointers((BONE*)bone->pFrameSibling);
	if(bone->pFrameFirstChild != NULL)SetupBoneMatrixPointers((BONE*)bone->pFrameFirstChild);

}

BONE* CSkinnedMesh::FindBone(char* szName)
{
	return (BONE*)D3DXFrameFind(m_pRootBone, szName);
}

void CSkinnedMesh::SetPose(D3DXMATRIX world, ID3DXAnimationController* animControl, float fTime)
{
	if(animControl != NULL)
		animControl->AdvanceTime(fTime, NULL);
	else
		m_pAnimControl->AdvanceTime(fTime, NULL);

	UpdateMatrices((BONE*)m_pRootBone, &world);
}

void CSkinnedMesh::SetAnimation(char* szName)
{
	ID3DXAnimationSet* anim = NULL;

	for(unsigned int nSet = 0; nSet < m_pAnimControl->GetMaxNumAnimationSets(); nSet++)
	{
		anim = NULL;
		m_pAnimControl->GetAnimationSet(nSet, &anim);

		if(anim != NULL)
		{
			if(strcmp(szName, anim->GetName()) == 0)
				m_pAnimControl->SetTrackAnimationSet(0, anim);
			anim->Release();
		}
	}
}

std::vector<std::string> CSkinnedMesh::GetAnimations()
{
	ID3DXAnimationSet* anim = NULL;
	std::vector<std::string> animations;

	for(unsigned int nSet = 0; nSet < m_pAnimControl->GetMaxNumAnimationSets(); nSet++)
	{
		anim = NULL;
		m_pAnimControl->GetAnimationSet(nSet, &anim);

		if(anim != NULL)
		{
			animations.push_back(anim->GetName());
			anim->Release();
		}

	}

	return animations;
}